import pygame
import random
import sys

# 初始化 pygame
pygame.init()

# 游戏窗口设置
WIDTH, HEIGHT = 640, 480
GRID_SIZE = 20
GRID_WIDTH = WIDTH // GRID_SIZE
GRID_HEIGHT = HEIGHT // GRID_SIZE

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLUE = (0, 0, 255)

# 蛇的初始设置
snake = [(GRID_WIDTH // 2, GRID_HEIGHT // 2)]
snake_direction = (1, 0)  # 初始方向：右
next_direction = snake_direction

# 食物设置
food = (random.randint(0, GRID_WIDTH - 1), random.randint(0, GRID_HEIGHT - 1))
while food in snake:
    food = (random.randint(0, GRID_WIDTH - 1), random.randint(0, GRID_HEIGHT - 1))

# 游戏状态
score = 0
game_over = False
clock = pygame.time.Clock()
FPS = 10

# 初始化窗口
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("贪吃蛇游戏")

# 字体设置
font = pygame.font.Font(None, 36)

# 游戏主循环
while not game_over:
    # 事件处理
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP and snake_direction != (0, 1):
                next_direction = (0, -1)
            elif event.key == pygame.K_DOWN and snake_direction != (0, -1):
                next_direction = (0, 1)
            elif event.key == pygame.K_LEFT and snake_direction != (1, 0):
                next_direction = (-1, 0)
            elif event.key == pygame.K_RIGHT and snake_direction != (-1, 0):
                next_direction = (1, 0)

    # 更新蛇的方向
    snake_direction = next_direction

    # 计算蛇头的新位置
    head_x, head_y = snake[0]
    dx, dy = snake_direction
    new_head = ((head_x + dx) % GRID_WIDTH, (head_y + dy) % GRID_HEIGHT)

    # 检查游戏结束条件
    if new_head in snake:
        game_over = True

    # 添加新蛇头
    snake.insert(0, new_head)

    # 检查是否吃到食物
    if new_head == food:
        score += 1
        # 生成新的食物
        food = (random.randint(0, GRID_WIDTH - 1), random.randint(0, GRID_HEIGHT - 1))
        while food in snake:
            food = (random.randint(0, GRID_WIDTH - 1), random.randint(0, GRID_HEIGHT - 1))
    else:
        # 移除蛇尾
        snake.pop()

    # 绘制背景
    screen.fill(BLACK)

    # 绘制蛇
    for segment in snake:
        x, y = segment
        pygame.draw.rect(screen, GREEN, (x * GRID_SIZE, y * GRID_SIZE, GRID_SIZE - 1, GRID_SIZE - 1))

    # 绘制食物
    food_x, food_y = food
    pygame.draw.rect(screen, RED, (food_x * GRID_SIZE, food_y * GRID_SIZE, GRID_SIZE - 1, GRID_SIZE - 1))

    # 绘制分数
    score_text = font.render(f"分数: {score}", True, WHITE)
    screen.blit(score_text, (10, 10))

    # 更新显示
    pygame.display.flip()

    # 控制游戏速度
    clock.tick(FPS)

# 游戏结束界面
game_over_font = pygame.font.Font(None, 72)
game_over_text = game_over_font.render("游戏结束", True, RED)
text_rect = game_over_text.get_rect(center=(WIDTH // 2, HEIGHT // 2))
screen.blit(game_over_text, text_rect)
pygame.display.flip()

# 等待退出
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                pygame.quit()
                sys.exit()